import os
import torch
import torchaudio
# from scipy.io import wavfile  # 看存音频部分的说明，用到这个包就导入即可

from model import TasNet


def separate(root, out_path):
    """ 对目标音频进行分离

    root: 混合音频的文件路径
    out_path: 分离后的音频保存的文件夹路径
    """

    # 加载模型参数
    net = TasNet()
    net.load_state_dict(torch.load("./save_net/epoch1-Conv-TasNet.pt", map_location=torch.device('cpu')))
    net.eval()

    # 加载音频数据并分离
    wf, sr = torchaudio.load(root)
    out = net(wf)[0]

    # 保存并输出分离的音频
    for i in range(2):
        ph = os.path.join(out_path, f"s{i}.wav")
        torchaudio.save(ph, out[i].reshape(1, -1), sr)
        """ 我在 windows 下通过 torchaudio.save 来保存时遇到问题，可以用以下代码替换 torchaudio.save
        # 需要安装 scipy 来导入 wavfile 包
        data = out[i].detach().numpy().reshape(-1, 1)
        wavfile.write(ph, sr, data)
        """
        print(f"save \'{ph}\'")


if __name__ == "__main__":
    separate("./data/TIMIT-mix/TEST/DR1/4541/4541-mix.wav", "./")
